As a general comment. I'm not sure I like this quick-and-dirty drop-in way of integrating Bmad-X elements. However, I think it makes sense to test them this way. In the long-term we should think about how this can be clean up better.

Some questions we should ask ourselves:

• Do we really want a separate bmadx module? This means that code is suddenly in a very different place than the other tracking code.
• Should there be a single Quadrupole class that has something like a tracking_method: Literal["linear", "bmadx"] argument?

Maybe slightly unrelated but I saw it in this PR:

• I think it's not very nice to have something like C = scipy.constants.speed_of_light at the top of a file. Ideally, one would just from scripy.constants import speed_of_light and use that. I think this might cause issues sometimes, when a torch.Tensor is strictly needed, but I cannot produce an example right now. Maybe we should add a cheetah.constants module that takes scripy.constnts and always returns torch.Tensors. But again, I'm not sure we actually need that ... just something to think about.

One small thing to keep in mind, bmadx tracking agrees to single precision with bmadx. I was expecting it to agree to double precision, but it is good enough for now.

@jp-ga, the question came up with @roussel-ryan, if this or #163 should be merged first. I think a main question in that was if the changes in #163 might make conversions in this PR redundant. Do you have an opinion on this?

I'm not sure if we should add now how Quadrupole.split handles the num_steps argument.

@jank324 I thought about this back when developing Bmad-X, and at the end I decided that it was better to keep the number of steps in the quad different and independent from the number of elements when splitting. For example, if you have a quad with num_steps=3 and you split the element in 5, the code should do 5 quadrupoles with L/5 length, each of which with num_steps=3 (doing 3 drift-kick-drift steps).

In the future I would like to implement fringe fields, so the splitting should have entrance fringe at the first element's entrance and exit fringe at the last element's exit, but we can leave this for a future PR.

I'm not sure if we should add now how Quadrupole.split handles the num_steps argument.

@jank324 I thought about this back when developing Bmad-X, and at the end I decided that it was better to keep the number of steps in the quad different and independent from the number of elements when splitting. For example, if you have a quad with num_steps=3 and you split the element in 5, the code should do 5 quadrupoles with L/5 length, each of which with num_steps=3 (doing 3 drift-kick-drift steps).

In the future I would like to implement fringe fields, so the splitting should have entrance fringe at the first element's entrance and exit fringe at the last element's exit, but we can leave this for a future PR.

@jp-ga So the behaviour you would like to have now is that each of the 5 child quadrupoles has 3 splits, just like their parent had 3 splits? In that case, don't we need to add the line highlighted below in the split method?

element = Quadrupole(
    torch.min(resolution, remaining),
    self.k1,
    misalignment=self.misalignment,
    tilt=self.tilt,
    num_steps=self.num_steps,   # <-- Add this line?
    dtype=self.length.dtype,
    device=self.length.device,
)

@jank324 I'm not sure if the test you added will raise the issue I ran into, I think my error had to do with the ellipsoid beam generation, but I'm not sure if that test will run into the same issue. I think it may be likely through since I remember that the different methods used the same code

I'm not sure if we should add now how Quadrupole.split handles the num_steps argument.

@jank324 I thought about this back when developing Bmad-X, and at the end I decided that it was better to keep the number of steps in the quad different and independent from the number of elements when splitting. For example, if you have a quad with num_steps=3 and you split the element in 5, the code should do 5 quadrupoles with L/5 length, each of which with num_steps=3 (doing 3 drift-kick-drift steps).
In the future I would like to implement fringe fields, so the splitting should have entrance fringe at the first element's entrance and exit fringe at the last element's exit, but we can leave this for a future PR.

@jp-ga So the behaviour you would like to have now is that each of the 5 child quadrupoles has 3 splits, just like their parent had 3 splits? In that case, don't we need to add the line highlighted below in the split method?

element = Quadrupole(
    torch.min(resolution, remaining),
    self.k1,
    misalignment=self.misalignment,
    tilt=self.tilt,
    num_steps=self.num_steps,   # <-- Add this line?
    dtype=self.length.dtype,
    device=self.length.device,
)

@jank324 Exactly, I will add it soon.

One final comment: Is it intended to keep the dev/ folder? If the notebook is useful, shall we move it to some other places, like a documentation page? Otherwise it's good to go!

I'm also not sure about the dev/ folder. I left it because I also sometimes write notebooks to develop features and I often worry that I will have to write them twice because I always delete them at some point once I'm done and I never commit them. dev/ would be a clean way to keep them around.

On the other hand, I've never actually felt like I needed to rewrite a notebook for some reason and we might not want to clutter up the repository and history like this.

So ... to document the decision: The dev/ folder remains for now, but the long-term goal should be to not have it. Interesting examples should instead, for example, appear in the documentation.